import zlib
from io import BytesIO

from scapy.all import *
from scapy.layers.inet import TCP, IP


def sanitize_filename(filename):
    # 使用正则表达式替换非法字符
    sanitized = re.sub(r'[<>:"/\\|?*\x00-\x1F]', '_', filename)
    # 限制文件名长度，避免操作系统的限制
    max_length = 255  # Windows 系统通常限制为255个字符
    return sanitized[:max_length]


def decode_content(content, encoding, brotli=None):
    # 根据content-Encoding 解码内容
    if encoding == 'gzip':
        with gzip.GzipFile(fileobj=BytesIO(content)) as gz:
            return gz.read()
    elif encoding == 'deflate':
        return zlib.decompress(content)
    elif encoding == 'brotli':
        return brotli.decompress(content)
    elif encoding == 'x-gzip':
        with gzip.GzipFile(fileobj=BytesIO(content)) as gz:
            return gz.read()
    elif encoding == 'compress':
        return zlib.decompress(content, wbits=zlib.MAX_WBITS | 16)  # 使用适当的 wbits 值
    else:
        return content  # 如果没有特定的编码，返回原始内容


def parse_pasv_command(command):
    """解析PASV命令以获取IP和端口"""
    if "Entering Passive Mode" in command:
        match = re.search(r'(\d+),(\d+),(\d+),(\d+),(\d+),(\d+)', command)
        if match:
            ip = '.'.join(match.groups()[:4])
            port = (int(match.group(5)) * 256) + int(match.group(6))
            return ip, port
    return None, None


def parse_port_command(command):
    """解析PORT命令以获取IP和端口"""
    if command.startswith("PORT"):
        parts = command.split()[1].split(',')
        ip = '.'.join(parts[:4])
        port = (int(parts[4]) * 256) + int(parts[5])
        return ip, port
    return None, None


def find_ftp_session_by_data_connection(src_ip, src_port, dest_ip, dest_port):
    # 查找是否存在匹配的数据连接
    for key, session in ftp_sessions.items():
        if session['current_data_connection'] in [(src_ip, src_port), (dest_ip, dest_port)]:
            return key, session
    return None, None


# 定义 MIME 类型和扩展名映射
mime_types = {
    'image/jpeg': 'jpg',
    'image/png': 'png',
    'image/gif': 'gif',
    'image/bmp': 'bmp',
    'text/html': 'html',
    'text/plain': 'txt',
    'text/javascript': 'js',
    'application/json': 'json',
    'application/javascript': 'js',
    'application/pdf': 'pdf',
    'application/xml': 'xml',
    'application/zip': 'zip',
    'application/octet-stream': 'bin'
}

session_http = {}
session_smtp = {}
# 用于存储每个FTP会话的状态
ftp_sessions = {}


def restore_protocol(pcap_file, output_dir):
    # 存储写入文件的完整路径
    file_list = []
    # 用来存储邮件的内容
    email_content = {}
    # 确保output_dir存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    with PcapReader(pcap_file) as reader:
        for data in reader:
            if data.haslayer(TCP) and data.haslayer(IP):
                src_ip = data[IP].src
                dest_ip = data[IP].dst
                s_port = data[TCP].sport
                d_port = data[TCP].dport
                flags = data[TCP].flags

                # 检查端口是否为80（HTTP）
                if s_port == 80 or d_port == 80:
                    # 处理相关逻辑
                    if s_port == 80:
                        key = (src_ip, dest_ip, s_port, d_port)
                        direction = 'client'  # server->client
                    elif d_port == 80:
                        key = (dest_ip, src_ip, d_port, s_port)
                        direction = 'server'  # client->server
                    if key not in session_http and flags == 'S':  # 刚刚建立
                        session_http[key] = {}
                    if Raw in data and flags not in ('FA', 'F') and key in session_http:  # 证明有负载且未结束
                        seq = data[TCP].seq
                        payload = bytes(data[Raw])
                        # 仅保留首次接收到的分片（去重）
                        if direction not in session_http[key]:
                            session_http[key][direction] = {}
                            # 检查是否存在相同序列号的数据包
                        if seq in session_http[key][direction]:
                            # 如果当前数据包的负载不为空且不同于已存储的数据
                            if payload and payload != session_http[key][direction][seq]:
                                # 将当前负载附加到已有负载
                                session_http[key][direction][seq] += payload
                                # print(f"检测到相同序列号的分片，已合并到序列号 {seq} 的数据。")
                        else:
                            # 仅在没有重复的情况下添加数据包
                            session_http[key][direction][seq] = payload
                    elif flags in ('FA', 'F') and key in session_http:
                        # 检查会话中是否存在双向流量的数据
                        file_counter = 0  # 用于区分文件名
                        for direction, seq_dict in session_http[key].items():
                            # 按照序号对负载进行排序，并确保乱序重传的到正确的重组
                            ordered_data = b''.join(
                                payload for seq, payload in sorted(seq_dict.items())
                            )
                            if direction == 'server':
                                # 从 HTTP 请求中提取文件路径
                                get_request_match = re.search(b'GET (/[^ ]+)', ordered_data)
                                if get_request_match:
                                    file_path = get_request_match.group(1).decode()
                                    request_file_name = os.path.basename(file_path)
                                    request_file_name = request_file_name.split('?')[0]  # 去掉查询参数
                                    request_file_name = sanitize_filename(request_file_name)

                                    if not request_file_name:
                                        request_file_name = f"extracted_file_{file_counter}.html"
                                else:
                                    request_file_name = None  # 未能从请求中提取文件名
                            elif direction == 'client':
                                # 分割出多个 HTTP 响应
                                responses = re.split(b'(HTTP/1.[01] \d{3} [^\r\n]+)', ordered_data)
                                for i in range(1, len(responses), 2):
                                    response_data = responses[i] + responses[i + 1] if i + 1 < len(responses) else \
                                        responses[i]
                                    # 获取 Content-Type 和 Content-Encoding
                                    content_type_match = re.search(b'Content-Type: ([\w/+-]+)', response_data)
                                    content_encoding_match = re.search(b'Content-Encoding: ([\w-]+)', response_data)
                                    content_type = content_type_match.group(1).decode() if content_type_match else 'application/octet-stream'
                                    encoding_type = content_encoding_match.group(1).decode() if content_encoding_match else None
                                    print(f"发现 Content-Type: {content_type}")
                                    # 获取文件扩展名
                                    file_ext = mime_types.get(content_type, 'dat')
                                    file1 = request_file_name if request_file_name else f"extracted_file_{file_counter}"
                                    # 检查文件名后缀
                                    if not file1.endswith(f'.{file_ext}'):
                                        file1 += f'.{file_ext}'  # 添加后缀
                                    # 确保文件名唯一
                                    base_file_name, ext = os.path.splitext(file1)  # 分离文件名和扩展名
                                    count = 1
                                    while os.path.exists(os.path.join(output_dir, file1)):
                                        file1 = f"{base_file_name}_{count}{ext}"  # 以 "a_1.html" 的格式生成文件名
                                        count += 1
                                    file_name = os.path.join(output_dir, file1)
                                    file_counter += 1
                                    # 提取响应体内容
                                    body_match = re.search(b'\r\n\r\n', response_data)
                                    if body_match:
                                        data_content = response_data[body_match.end():]
                                        # 解码内容
                                        data_content = decode_content(data_content, encoding_type)
                                        # 根据 MIME 类型处理内容
                                        if content_type.startswith('image/'):
                                            with open(file_name, "wb") as file:
                                                file.write(data_content)
                                            file_list.append(file_name)
                                            print(f"图片已保存为 {file_name}")
                                        elif content_type in ('text/html', 'text/plain', 'text/javascript'):
                                            with open(file_name, "w", encoding="utf-8") as file:
                                                file.write(data_content.decode(errors="ignore"))
                                            file_list.append(file_name)
                                            print(f"文本文件已保存为 {file_name}")
                                        elif content_type == 'application/json':
                                            with open(file_name, "w", encoding="utf-8") as file:
                                                file.write(data_content.decode(errors="ignore"))
                                            file_list.append(file_name)
                                            print(f"JSON 文件已保存为 {file_name}")
                                        elif content_type == 'application/pdf':
                                            with open(file_name, "wb") as file:
                                                file.write(data_content)
                                            file_list.append(file_name)
                                            print(f"PDF 文件已保存为 {file_name}")
                                        elif content_type == 'application/zip':
                                            with open(file_name, "wb") as file:
                                                file.write(data_content)
                                            file_list.append(file_name)
                                            print(f"ZIP 文件已保存为 {file_name}")
                                        elif content_type == 'application/octet-stream':
                                            with open(file_name, "wb") as file:
                                                file.write(data_content)
                                            file_list.append(file_name)
                                            print(f"二进制文件已保存为 {file_name}")
                                        else:
                                            # 处理未知类型
                                            with open(file_name, "wb") as file:
                                                file.write(data_content)
                                            file_list.append(file_name)
                                            print(f"未知类型文件已保存为 {file_name}，Content-Type: {content_type}")
                        del (session_http[key])
                # 检查端口是否为25（SMTP）
                elif s_port == 25 or d_port == 25:
                    if s_port == 25:
                        key = (src_ip, dest_ip, s_port, d_port)
                        direction = 'client'  # server->client
                    elif d_port == 25:
                        key = (dest_ip, src_ip, d_port, s_port)
                        direction = 'server'  # client->server
                    # 初始化会话
                    if key not in session_smtp and flags == 'S':
                        session_smtp[key] = {}
                    if Raw in data and flags not in ('FA', 'F') and key in session_smtp:  # 证明有负载且未结束
                        seq = data[TCP].seq
                        payload = bytes(data[Raw])
                        # payload_len = len(payload)
                        # 确保方向字典已初始化
                        if direction not in session_smtp[key]:
                            session_smtp[key][direction] = {}
                        for i, byte in enumerate(payload):
                            session_smtp[key][direction][seq + i] = byte
                    elif flags in ('FA', 'F') and key in session_smtp:
                        # 检查会话中是否存在双向流量的数据
                        for direction, seq_dict in session_smtp[key].items():
                            # 按序列号排序去处理乱序问题
                            sorted_data = sorted(seq_dict.items())
                            full_data = b''.join([bytes([byte]) for _, byte in sorted_data])
                            lines = full_data.split(b'\r\n')
                            in_email = False
                            current_email = []
                            for line in lines:
                                if line.startswith(b'DATA'):
                                    in_email = True
                                    continue
                                if line == b'.':
                                    # 结束当前邮 件
                                    in_email = False
                                    # 将邮件内容保存为.eml文件
                                    eml_filename = os.path.join(output_dir, f'email_{d_port}.eml')
                                    with open(eml_filename, 'wb') as eml_file:
                                        eml_file.write(b'\r\n'.join(current_email) + b'\r\n.')
                                    file_list.append(eml_filename)
                                    print(f"邮件已保存为 {eml_filename}")
                                    email_content[eml_filename] = current_email
                                    current_email = []
                                    continue
                                if in_email:
                                    current_email.append(line)
                        # return eml_filename
                        del session_smtp[key]
                elif s_port == 21 or d_port == 21:
                    if s_port == 21:
                        key = (src_ip, dest_ip, s_port, d_port)
                    elif d_port == 21:
                        key = (dest_ip, src_ip, d_port, s_port)
                    # 检查是否已有该会话的状态，并且是控制连接刚刚SYN的阶段
                    if key not in ftp_sessions:
                        # 初始化会话的状态
                        ftp_sessions[key] = {
                            'mode': 1,  # 初始为控制连接模式
                            'current_data_connection': None,
                            'data_command': "",
                            'data_packets': defaultdict(list)
                        }
                    session = ftp_sessions[key]

                    # 处理FTP的控制连接
                    if session['mode'] == 1:
                        if data[TCP].payload:
                            payload = bytes(data[TCP].payload.load).decode('utf-8', errors='ignore')

                            if payload.startswith("PORT"):
                                session['current_data_connection'] = parse_port_command(payload)
                            elif "Entering Passive Mode" in payload:
                                session['current_data_connection'] = parse_pasv_command(payload)
                            elif payload.startswith(("LIST", "RETR", "DELE", "STOR", "NLST", "MLSD")) and session[
                                'current_data_connection']:
                                session['mode'] = 2
                                session['data_command'] = payload.strip()
                # 处理FTP数据连接
                if find_ftp_session_by_data_connection(src_ip, s_port, dest_ip, d_port):
                    key, session = find_ftp_session_by_data_connection(src_ip, s_port, dest_ip, d_port)
                    # 如果找到了匹配的会话且处于数据模式
                    if session and session['mode'] == 2:
                        # 记录数据包
                        if Raw in data:
                            seq_num = data[TCP].seq
                            session['data_packets'][seq_num] = data[Raw].load
                        # 检查数据连接是否完成
                        if flags in ('FA', 'F'):
                            sorted_data = bytearray()
                            for seq_num in sorted(session['data_packets'].keys()):
                                sorted_data.extend(session['data_packets'][seq_num])

                            command = session['data_command']
                            if command.startswith("RETR"):
                                filename = command.split()[1]
                                output_filename = f"RETR_{filename}"
                                print(
                                    f"Preparing to save data for RETR command as {output_filename}")  # 调试：打印RETR文件名
                            elif command.startswith("STOR"):
                                filename = command.split()[1]
                                output_filename = f"STOR_{filename}"
                                print(f"Preparing to save data for STOR command as {output_filename}")
                            elif command.startswith(("LIST", "NLST", "MLSD")):
                                output_filename = f"LIST_{d_port}.txt"
                                print(f"Preparing to save data for LIST command as {output_filename}")
                            else:
                                output_filename = "ftp_file_restore.txt"
                            file_path = os.path.join(output_dir, output_filename)
                            try:
                                with open(file_path, 'wb') as f:
                                    f.write(sorted_data)
                                file_list.append(file_path)
                                print(f"Data for {command} written to {file_path}")
                            except Exception as e:
                                print(f"Error writing to file {file_path}: {e}")
                            ftp_sessions.pop(key, None)
    return file_list


if __name__ == '__main__':
    file_list = restore_protocol("D:\\PycharmProjects\\Protocol_parse\\pcap_file\\HTTP\\http_gzip.cap",
                                 "D:\\PycharmProjects\\Protocol_parse\\restore\\restore")
    for i in file_list:
        print(i)
